import numpy as np
from python_ai.common.xcommon import *

np.set_printoptions(edgeitems=200)

m, n = 4, 4  # ori 13, 13
box_per_cell = 3 # ori 5
n_cls = 5  # ori 20

sep('confidence')
n_el = m*n*box_per_cell
print('n_el', n_el)
confidence = np.arange(m*n*box_per_cell).reshape([m, n, box_per_cell])
print(confidence.shape)

sep('expand_dims')
confidence = np.expand_dims(confidence, 3)
print(confidence.shape)

sep('tile')
confidence = np.tile(confidence, (1, 1, 1, n_cls))
print(confidence.shape)

sep('nonzero')
classes = confidence / n_el
filter_probs = np.array(classes >= 0.5, dtype='bool')  # (m, n, box_per_cell, n_cls)
print(filter_probs.shape)
print(filter_probs)
filter_index = np.nonzero(filter_probs)
for el in filter_index:
    print(el.shape)
print(filter_index)

